#!/usr/bin/env python3

"""
    chassis_db_init
    Chassis information update tool for SONiC
    This tool runs one time at the launch of the platform monitor in order to populate STATE_DB with chassis information such as model, serial number, and revision.
"""

try:
    import os
    import sys

    from sonic_py_common import daemon_base, logger

    # If unit testing is occurring, mock swsscommon and module_base
    if os.getenv("CHASSIS_DB_INIT_UNIT_TESTING") == "1":
        from tests import mock_swsscommon as swsscommon
    else:
        from swsscommon import swsscommon
except ImportError as e:
    raise ImportError(str(e) + " - required module not found")

#
# Constants
#

SYSLOG_IDENTIFIER = "chassis_db_init"

CHASSIS_INFO_TABLE = 'CHASSIS_INFO'
CHASSIS_INFO_KEY_TEMPLATE = 'chassis {}'
CHASSIS_INFO_CARD_NUM_FIELD = 'module_num'
CHASSIS_INFO_SERIAL_FIELD = 'serial'
CHASSIS_INFO_MODEL_FIELD = 'model'
CHASSIS_INFO_REV_FIELD = 'revision'

CHASSIS_LOAD_ERROR = 1

NOT_AVAILABLE = 'N/A'

#
# Helper functions =============================================================
#

# try get information from platform API and return a default value if caught NotImplementedError


def try_get(callback, *args, **kwargs):
    """
    Handy function to invoke the callback and catch NotImplementedError
    :param callback: Callback to be invoked
    :param args: Arguments to be passed to callback
    :param kwargs: Default return value if exception occur
    :return: Default return value if exception occur else return value of the callback
    """
    default = kwargs.get('default', NOT_AVAILABLE)
    try:
        ret = callback(*args)
        if ret is None:
            ret = default
    except NotImplementedError:
        ret = default

    return ret

#
# Functions
# 

def provision_db(platform_chassis, log):
    # Init state db connection
    state_db = daemon_base.db_connect("STATE_DB")
    chassis_table = swsscommon.Table(state_db, CHASSIS_INFO_TABLE)

    # Populate DB with chassis hardware info 
    fvs = swsscommon.FieldValuePairs([
                                        (CHASSIS_INFO_SERIAL_FIELD, try_get(platform_chassis.get_serial)),
                                        (CHASSIS_INFO_MODEL_FIELD, try_get(platform_chassis.get_model)),
                                        (CHASSIS_INFO_REV_FIELD, try_get(platform_chassis.get_revision))
                                    ])
    chassis_table.set(CHASSIS_INFO_KEY_TEMPLATE.format(1), fvs)
    log.log_info("STATE_DB provisioned with chassis info.")

    return chassis_table


#
# Main
#

def main():
    log = logger.Logger(SYSLOG_IDENTIFIER)
    log.log_info("Provisioning Database with Chassis Info...")

    # Load platform api class
    try:
        import sonic_platform.platform
        platform_chassis = sonic_platform.platform.Platform().get_chassis()
    except Exception as e:
        log.log_error("Failed to load chassis due to {}".format(repr(e)))
        sys.exit(CHASSIS_LOAD_ERROR)

    provision_db(platform_chassis, log)

if __name__ == '__main__':
    main()
